import base64
from Crypto import Random
from Crypto.Hash import SHA256
from Crypto.PublicKey import RSA
from Crypto.Signature import PKCS1_v1_5 as Signature_pkcs1_v1_5


def get_key():
    """
    生成公私钥
    :return: 公、私钥base64编码字符串
    """
    # 获取一个伪随机数生成器
    random_generator = Random.new().read
    # 获取一个rsa算法对应的密钥对生成器实例
    rsa = RSA.generate(2048, random_generator)
    # 生成私钥并保存
    private_pem = rsa.exportKey()
    # 生成公钥并保存
    public_pem = rsa.publickey().exportKey()
    return str(base64.b64encode(public_pem), encoding="utf8"), str(base64.b64encode(private_pem), encoding="utf8")


def gen_sign(private_key, unsign_data):
    """
    私钥签名
    :param private_key: 私钥base64编码字符串
    :param unsign_data: 待签名数据
    :return: 签名字符串
    """
    rsaKey = RSA.importKey(base64.b64decode(bytes(private_key, encoding="utf8")))
    signer = Signature_pkcs1_v1_5.new(rsaKey)
    digest = SHA256.new()
    digest.update(unsign_data.encode('utf8'))
    sign = signer.sign(digest)
    signature = str(base64.b64encode(sign), encoding='utf8')
    return signature


def verify_sign(pubkey, data, sign):
    """

    :param pubkey: 公钥base64编码字符串
    :param data: 原数据
    :param sign: 签名字符串
    :return: 验签结果true，false
    """
    rsaKey = RSA.importKey(base64.b64decode(bytes(pubkey, encoding="utf8")))
    verifier = Signature_pkcs1_v1_5.new(rsaKey)
    digest = SHA256.new()
    digest.update(data.encode('utf8'))
    is_verify = verifier.verify(digest, base64.b64decode(sign))
    return is_verify


if __name__ == '__main__':
    data = 'boblee'
    pub, prv = get_key()
    sign = gen_sign(prv, data)
    result = verify_sign(pub, data, sign)
